from multiprocessing import Process
import os
import sys
from threading import Thread
from time import sleep
import psutil

class StatusThread(Thread):
    tname = ""
    pid = 0
    end_flag = False
    save_dir = ""

    def __init__(self, pid, tname, save_dir):
        super().__init__()
        self.tname = tname
        self.pid = pid
        self.save_dir = save_dir

    def end(self):
        self.end_flag = True

    def run(self):
        p = psutil.Process(self.pid)
        cpu_avg = 0
        mem_avg = 0
        cpu_max = 0
        mem_max = 0
        cpu_min = float('inf')
        mem_min = float('inf')
        times = 1
        f = open(self.save_dir + self.tname + "_stat.txt", "w")
        while p.is_running() and not self.end_flag:
            try:
                cpu = p.cpu_percent()
                mem = p.memory_full_info().rss / 1024. / 1024.
                cpu_max = max(cpu_max, cpu)
                mem_max = max(mem_max, mem)
                cpu_min = min(cpu_min, cpu)
                mem_min = min(mem_min, mem)
                cpu_avg += cpu
                mem_avg += mem
                f.write("%f %f\n"%(cpu, mem))
                times += 1
                sleep(0.1)
            except:
                break
        print(self.tname + " Status:")
        print("Avg: cpu: %f\t mem: %f"%(cpu_avg/times, mem_avg/times))
        print("Max: cpu: %f\t mem: %f"%(cpu_max, mem_max))
        print("Min: cpu: %f\t mem: %f"%(cpu_min, mem_min))

        f = open(self.save_dir + self.tname + "_stat_sum.txt", "w")
        f.write("Avg: cpu: %f\t mem: %f\n"%(cpu_avg/times, mem_avg/times))
        f.write("Max: cpu: %f\t mem: %f\n"%(cpu_max, mem_max))
        f.write("Min: cpu: %f\t mem: %f\n"%(cpu_min, mem_min))


def find_process(min_pid, process_name):
    pids = psutil.pids()
    target = 0
    for pid in pids:
        p = psutil.Process(pid)
        if p.pid >= min_pid and p.name() == process_name:
            target = pid
            break
    return target

if __name__ == '__main__':
    pid = sys.argv[1]
    StatusThread(int(pid), "tmp_status.txt").start()
